/*
 * Copyright 2023 Huawei Cloud Computing Technology Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.huawei.cloudphone.virtualdevice.camera;

import static android.media.MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar;
import static android.media.MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
import static android.os.Build.VERSION_CODES.M;

import android.graphics.ImageFormat;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaCodecList;
import android.media.MediaFormat;
import android.os.Build;
import android.os.Bundle;
import android.util.Log;

import java.io.IOException;
import java.nio.ByteBuffer;

public class AvcEncoder {
    private static final String TAG = "AvcEncoder";
    private static final String MIME_TYPE = "video/avc";
    private static final int TIMEOUT_USC = 1000;
    private static final int MICROSECONDS_PER_SECOND = 1000000;
    private static final int I_FRAME_INTERVAL = 150;

    private static int colorFormat = COLOR_FormatYUV420Planar;
    private MediaCodec mediaCodec;
    private int mWidth;
    private int mHeight;
    private int mFrameRate;
    private long genIndex;
    private byte[] configByte;
    private boolean isRequestKeyFrame = false;
    private AvcEncoderDataHandler dataHandler;

    public void flush() {
        if (mediaCodec != null) {
            mediaCodec.flush();
        }
    }

    public void generateKeyFrame() {
        if (Build.VERSION.SDK_INT >= M && mediaCodec != null) {
            Bundle params = new Bundle();
            params.putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 0);
            mediaCodec.setParameters(params);
        }
    }

    // 支持MediaCodec动态设置上行传输码率
    public void updateDynamicBitrate(int newBitrate) {
        if (Build.VERSION.SDK_INT >= M && mediaCodec != null) {
            Bundle params = new Bundle();
            params.putInt(MediaCodec.PARAMETER_KEY_VIDEO_BITRATE, newBitrate);
            mediaCodec.setParameters(params);
        }
    }

    public AvcEncoder(int width, int height, int frameRate, int bitrate, AvcEncoderDataHandler dataHandler) {
        Log.i(TAG, "Create AvcEncoder, width = " + width + ", height = " + height +
                ", frameRate = " + frameRate + ", bitrate = " + bitrate);
        this.dataHandler = dataHandler;
        this.mWidth = width;
        this.mHeight = height;
        this.mFrameRate = frameRate;
        this.genIndex = 0;
        MediaFormat mediaFormat = MediaFormat.createVideoFormat(MIME_TYPE, width, height);
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, colorFormat);
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, I_FRAME_INTERVAL);
        mediaFormat.setInteger(MediaFormat.KEY_BITRATE_MODE, MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_CBR);

        try {
            mediaCodec = MediaCodec.createEncoderByType(MIME_TYPE);
        } catch (IOException e) {
            Log.e(TAG, "AvcEncoder create encoder by type failed. ", e);
        }
        mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        mediaCodec.start();
    }

    public void stopEncoder() {
        Log.i(TAG, "Stop encoder.");
        try {
            mediaCodec.stop();
            mediaCodec.release();
        } catch (IllegalStateException e) {
            Log.e(TAG, "AvcEncoder stop Encoder failed. ", e);
        }
    }

    public void offerEncoder(byte[] frame, byte[] encodeBuff) {
        if ((frame == null) || (encodeBuff == null)) {
            Log.e(TAG, "Offer Encoder input param invalid. ");
            return;
        }
        inputBuffer(frame);
        int outBufLen = 0;
        while (true) {
            if (genIndex % I_FRAME_INTERVAL == 0 && !isRequestKeyFrame) {
                generateKeyFrame();
                isRequestKeyFrame = true;
            }
            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            int outBufId = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USC);
            if (outBufId >= 0) {
                ByteBuffer outBuf = mediaCodec.getOutputBuffer(outBufId);
                byte[] h264Buf = new byte[bufferInfo.size];
                if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_CODEC_CONFIG) {
                    configByte = new byte[bufferInfo.size];
                    outBuf.get(configByte);
                } else if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_KEY_FRAME) {
                    outBuf.get(h264Buf);
                    System.arraycopy(configByte, 0, encodeBuff, 0, configByte.length);
                    System.arraycopy(h264Buf, 0, encodeBuff, configByte.length, h264Buf.length);
                    outBufLen = bufferInfo.size + configByte.length;
                    dataHandler.handleData(encodeBuff, outBufLen);
                } else {
                    outBuf.get(h264Buf);
                    System.arraycopy(h264Buf, 0, encodeBuff, 0, h264Buf.length);
                    outBufLen = bufferInfo.size;
                    dataHandler.handleData(encodeBuff, outBufLen);
                }
                mediaCodec.releaseOutputBuffer(outBufId, false);
                isRequestKeyFrame = false;
            } else if (outBufId == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                Log.e(TAG, "Output format changed : " + outBufId);
                return;
            } else if (outBufId == MediaCodec.INFO_TRY_AGAIN_LATER) {
                return;
            } else {
                Log.e(TAG, "Unknown outBufId : " + outBufId);
                return;
            }
        }
    }

    private void inputBuffer(byte[] frame) {
        long pts = computePresentationTime(genIndex);
        genIndex++;
        byte[] yuv420sp = new byte[mWidth * mHeight * 3 / 2];
        if (colorFormat == COLOR_FormatYUV420Planar) {
            swapYV12ToI420(frame, yuv420sp, mWidth, mHeight);
        }else if (colorFormat == COLOR_FormatYUV420SemiPlanar) {
            swapNV21ToNV12(frame, yuv420sp, mWidth, mHeight);
        }
        int inBuffId = mediaCodec.dequeueInputBuffer(TIMEOUT_USC);
        if (inBuffId < 0) {
            Log.e(TAG, "Dequeue input buffer error, inBuffId = " + inBuffId);
            return;
        }
        ByteBuffer inBuff = mediaCodec.getInputBuffer(inBuffId);
        inBuff.clear();
        inBuff.put(yuv420sp);
        mediaCodec.queueInputBuffer(inBuffId, 0, yuv420sp.length, pts, 0);
    }

    private void swapYV12ToI420(byte[] yv12Data, byte[] i420Data, int width, int height) {
        if (yv12Data == null || i420Data == null) {
            return;
        }
        int frameSize = width * height;
        System.arraycopy(yv12Data, 0, i420Data, 0, frameSize);
        System.arraycopy(yv12Data, frameSize + frameSize / 4, i420Data, frameSize ,frameSize / 4);
        System.arraycopy(yv12Data, frameSize, i420Data, frameSize + frameSize / 4, frameSize / 4);
    }

    private void swapNV21ToNV12(byte[] nv21Data, byte[] nv12Data, int width, int height) {
        if (nv21Data == null || nv12Data == null) {
            return;
        }
        int frameSize = width * height;
        System.arraycopy(nv21Data, 0, nv12Data, 0, frameSize);
        for (int i = 0; i < frameSize / 2; i += 2) {
            nv12Data[frameSize + i - 1] = nv21Data[i + frameSize];
        }
        for (int i = 0; i < frameSize / 2; i += 2) {
            nv12Data[frameSize + i] = nv21Data[i + frameSize - 1];
        }
    }

    private long computePresentationTime(long frameIndex) {
        return frameIndex * MICROSECONDS_PER_SECOND / mFrameRate;
    }

    public static int getImageFormat() {
        return colorFormat == COLOR_FormatYUV420Planar ? ImageFormat.YV12 : ImageFormat.NV21;
    }

    public static boolean isSupportH264() {
        MediaCodecInfo codecInfo = selectCodec(MIME_TYPE);
        if (codecInfo == null) {
            Log.e(TAG, "Couldn't find a codec for " + MIME_TYPE);
            return false;
        }
        return selectColorFormat(codecInfo, MIME_TYPE);
    }

    private static MediaCodecInfo selectCodec(String mimeType) {
        for (int i = 0; i < MediaCodecList.getCodecCount(); i++) {
            MediaCodecInfo codecInfo = MediaCodecList.getCodecInfoAt(i);
            if(!codecInfo.isEncoder()) {
                continue;
            }
            String[] types = codecInfo.getSupportedTypes();
            for (int j = 0; j < types.length; j++) {
                if (types[j].equalsIgnoreCase(mimeType)) {
                    return codecInfo;
                }
            }
        }
        return null;
    }

    private static boolean selectColorFormat(MediaCodecInfo codecInfo, String mimeType) {
        MediaCodecInfo.CodecCapabilities capabilities;
        try {
            capabilities = codecInfo.getCapabilitiesForType(mimeType);
        } catch (IllegalArgumentException e) {
            Log.e(TAG, "Couldn't get capabilities for " + MIME_TYPE);
            return false;
        }
        for (int i = 0; i < capabilities.colorFormats.length; i++) {
            int colorFormat = capabilities.colorFormats[i];
            if (colorFormat == COLOR_FormatYUV420SemiPlanar) {
                AvcEncoder.colorFormat = COLOR_FormatYUV420SemiPlanar;
                return true;
            } else if (colorFormat == COLOR_FormatYUV420Planar) {
                AvcEncoder.colorFormat = COLOR_FormatYUV420Planar;
                return true;
            }
        }
        Log.e(TAG, "Couldn't get a color format for " + codecInfo.getName() + "/" + mimeType);
        return false;
    }
}
